# Copyright (c) 2024 SwanLink (Jiangsu) Technology Development Co., LTD.
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

#/bin/sh

# apt install iputils-ping

set -e

is_debug_mode=0
is_unc_path=0
cur_dir=$(readlink -f .)
echo "cur_dir: $cur_dir"
script_name=$(realpath --relative-to=$cur_dir "$(readlink -f $0)")

function debug_print
{
    if [ $is_debug_mode -eq 1 ]; then
        echo "$*"
    fi
}

function check_temp_dir
{
    if [[ ! $1 =~ ^/tmp/tmp\.[a-zA-Z0-9]+$ ]]; then
        echo "ERR: check tmp dir fail: $1"
        exit 1
    fi
}

function usage()
{
    args1=$1
    if [[ ! -z $args1 ]]; then
        echo "Err: invalid argument: $args1"
    fi
    cat <<-EOF
    help:
    Usage: ${script_name} [ --local_pc | --ip <target_pc_ip> ] [[ --image_path | -i <image_path> ] | [ -t mount_type ] [ -o unc_path_mount_options ] [ --code_path | -c <code_path> ] [ --disable_reboot | -r ] [ --quiet | -q ] [ --debug | -d ]

    --local_pc | -l                 install local pc
    --ip ip                         target pc ip
    --image_path | -i image_path    image path, storing files such as system.img, vendor.img, bzImage, and etc.
                                    !!! ATTENTION !!! if image_path is a windows unc path, please quoted with single quote.
    -t mount_type                   mount_type, e.g. cifs
    -o unc_path_mount_options       mount options for unc path mount, e.g.
                                    1) windows path: e.g. -o username=root,password=123456
                                    2) linux path: There is no adaptation
    --code_path | -c code_path      developer code path, which is repo home.
    --disable_reboot | -r           do not reboot when image installed
    --quiet | -q                    quiet execute, no answer
    --debug | -d                    debug mode
    example:
    sudo bash ${script_name} --local_pc --code_path /home/code
    sudo bash ${script_name} --ip 10.10.10.10 --image_path '\\\\10.10.10.1\\daily_build\\image' -o username=root,password=123456
    sudo bash ${script_name} --ip 10.10.10.10 --image_path /home/code/out/jiuxiao/packages/phone/images
    sudo bash ${script_name} --ip 10.10.10.10 --code_path /home/code
EOF
    exit 0
}

disk_part_info=
function get_disk_part_info()
{
    is_install_from_local=$1
    if [[ $is_install_from_local -eq 0 ]]; then
        ip=$2
        disk_part_info=$(ssh root@$ip "lsblk --tree -o PATH,SIZE,TYPE,partlabel,fstype")
        if [[ $? -ne 0 ]]; then
            echo "ERR: please check target machine ip ($ip), it can reachable?"
            exit 1
        fi
    else
        disk_part_info=$(lsblk --tree -o PATH,SIZE,TYPE,partlabel,fstype)
        if [[ $? -ne 0 ]]; then
            echo "ERR: please check target machine ip ($ip), it can reachable?"
            exit 1
        fi
    fi
    debug_print "disk_part_info:\n$disk_part_info"
    if [[ -z "$disk_part_info" ]]; then
        echo "ERR: disk part info is empty"
        exit 2
    fi
}

function get_part_path()
{
    partlabel=$1
    ret=$(echo "$disk_part_info" | grep -v "vfat" | grep -v "Basic" | grep ${partlabel} | awk '{print $1;}')
    if [[ -z "$ret" ]]; then
        echo "ERR: can not get partlabel of $partlabel"
        exit 2
    fi
    echo "$ret"
    exit 0
}

is_local_pc=0
target_pc_ip=
image_path=
mount_type=
unc_path_mount_options=
code_path=
flag_reboot=1
flag_quiet=0

set +e
while [ $# -gt 0 ]; do
  debug_print "current parse args: $1"
  case $1 in
    --help | -h)
      usage
      ;;
    --local_pc | -l) is_local_pc=1; shift; ;;
    --ip) shift; target_pc_ip=${1}; shift; ;;
    --image_path | -i) shift; image_path=${1}; shift; ;;
    --mount_type | -t) shift; mount_type=${1}; shift; ;;
    --unc_path_mount_options | -o) shift; unc_path_mount_options=${1}; shift; ;;
    --code_path | -c) shift; code_path=${1}; shift; ;;
    --disable_reboot | -r) flag_reboot=0; shift; ;;
    --quiet | -q) flag_quiet=1; shift; ;;
    --debug | -d) is_debug_mode=1; shift; ;;
    *) usage $1; exit 0; shift; ;;
  esac
done
set -e

echo $is_local_pc
if [ $is_local_pc -eq 0 ] && [ -z "$target_pc_ip" ]; then
    echo -e "    ERROR: Please provide local_pc or target_pc_ip.\n"
    usage
    exit 1
fi

if [ $is_local_pc -ne 0 ] && [ ! -z "$target_pc_ip" ]; then
    echo -e "    ERROR: conflict, reson: both provide local_pc and target_pc_ip.\n"
    usage
    exit 1
fi

if [ ! -z "$target_pc_ip" ]; then
    is_local_pc=0
    if [ ! -f ~/.ssh/id_rsa.pub ]; then
        ssh-keygen -t rsa -N '' -f ~/.ssh/id_rsa -q
    fi
    ssh-copy-id -i ~/.ssh/id_rsa.pub root@$target_pc_ip
fi

if [ -z "$image_path" ] && [ -z "$code_path" ]; then
    echo -e " ERROR: Unknown image_path:$image_path, nor code_path:$code_path\n"
    usage
    exit 1
fi
if [ ! -z "$image_path" ] && [ ! -z "$code_path" ]; then
    echo -e " ERROR: Do not provide both image_path and code_path\n"
    usage
    exit 1
fi
temp_image_path_for_unc_mount=
if [ ! -z "$code_path" ]; then
    image_path=$(realpath $(realpath $code_path)/out/jiuxiao/packages/phone/images/)
    if [ ! -d $image_path ]; then
        echo "There is no image in dir: code_path($code_path), please check $image_path"
        usage
        exit 1
    fi
else
    debug_print "raw image_path: $image_path"
    debug_print "raw image_path left 2 char: ${image_path:0:2}"
    if [ ${image_path:0:2} == '\\' ]; then
        mount_type="cifs"
        normalized_image_path=$(echo ${image_path} | sed 's=\\=\/=g')
        debug_print "normalized image_path: $normalized_image_path"
        image_path=$(mktemp -d)
        temp_image_path_for_unc_mount=$image_path
    fi
    if [ ! -z $mount_type ]; then
        mount_cmd=
        if [ "$mount_type" == "cifs" ]; then
            mount_cmd="mount -t cifs"
        else
            mount_cmd="mount "
        fi
        if [ ! -z "$unc_path_mount_options" ]; then
            mount_cmd="$mount_cmd -o $unc_path_mount_options"
        fi
        mount_cmd="$mount_cmd ${normalized_image_path} $image_path"
        echo "mount_cmd: $mount_cmd"
        $($mount_cmd)
        is_unc_path=1
    fi

    image_path=$(realpath $image_path)
    if [[ ! -d ${image_path} ]]; then
        echo "image_path($image_path), not exists"
        usage
        exit 1
    fi
fi

[[ -f ${image_path}/system.img ]]
[[ -f ${image_path}/vendor.img ]]
[[ -f ${image_path}/ramdisk.img ]]
[[ -f ${image_path}/bzImage ]]
[[ -f ${image_path}/sys_prod.img ]]
[[ -f ${image_path}/chip_prod.img ]]

get_disk_part_info $is_local_pc $target_pc_ip
part_system=$(get_part_path system)
part_vendor=$(get_part_path vendor)
part_data=$(get_part_path data)
part_misc=$(get_part_path misc)
part_sys_prod=$(get_part_path sys_prod)
part_chip_prod=$(get_part_path chip_prod)

echo "!!! PLEASE CONFIRM PARTION !!!"
printf "%20s====%-75s\n" $(printf '=%.0s' {1..20}) $(printf '=%.0s' {1..75})
printf "%20s    %-75s\n" ITEM VALUE
printf "%20s====%-75s\n" $(printf '=%.0s' {1..20}) $(printf '=%.0s' {1..75})
printf "%20s    %-75s\n" is_debug_mode  "$([ $is_debug_mode -eq 1 ] && echo "true" || echo "false")"
printf "%20s    %-75s\n" is_local_pc    "$([ $is_local_pc -eq 1 ] && echo "true" || echo "false")"
printf "%20s    %-75s\n" target_pc_ip   "$([ ! -z $target_pc_ip ] && echo "$target_pc_ip" || echo "NULL")"
printf "%20s    %-75s\n" flag_reboot    "$([ $flag_reboot -eq 1 ] && echo "true" || echo "false")"
printf "%20s    %-75s\n" quiet          "$([ $flag_quiet -eq 1 ] && echo "true" || echo "false")"
printf "%20s    %-75s\n" image_path     "$image_path"
printf "%20s    %-75s\n" system         "$part_system"
printf "%20s    %-75s\n" vendor         "$part_vendor"
printf "%20s    %-75s\n" data           "$part_data"
printf "%20s    %-75s\n" misc           "$part_misc"
printf "%20s    %-75s\n" sys_prod       "$part_sys_prod"
printf "%20s    %-75s\n" chip_prod      "$part_chip_prod"
printf "%20s====%-75s\n" $(printf '=%.0s' {1..20}) $(printf '=%.0s' {1..75})

if [[ $flag_quiet -eq 1 ]]; then
    echo "Quiet mode"
else
    echo "Do you wish to continue install ?"
    select yn in "Yes" "No"; do
        case $yn in
            Yes ) break;;
            No ) exit;;
        esac
    done
fi

function gen_fstab_file
{
    fstab_tmp_file=$(mktemp)
    cat > $fstab_tmp_file <<-EOF
# fstab file.
#<src>                       <mnt_point>   <type>   <mnt_flags and options>         <fs_mgr_flags>
##part_system##              /usr          ext4     ro,barrier=1                    wait,required
##part_vendor##              /vendor       ext4     ro,barrier=1                    wait,required
##part_data##              /data         ext4     discard,noatime,nosuid,nodev,usrquota  wait,check,quota
##misc##              /misc         none     none                            wait,required
##part_sys_prod##              /sys_prod     ext4     ro,barrier=1                    wait
##part_chip_prod##              /chip_prod    ext4     ro,barrier=1                    wait
EOF

    sed -i "s=##part_system##=${part_system/\/dev\//\/dev\/block\/}=g"      $fstab_tmp_file
    sed -i "s=##part_vendor##=${part_vendor/\/dev\//\/dev\/block\/}=g"      $fstab_tmp_file
    sed -i "s=##part_data##=${part_data/\/dev\//\/dev\/block\/}=g"        $fstab_tmp_file
    sed -i "s=##misc##=${part_misc/\/dev\//\/dev\/block\/}=g"        $fstab_tmp_file
    sed -i "s=##part_sys_prod##=${part_sys_prod/\/dev\//\/dev\/block\/}=g"    $fstab_tmp_file
    sed -i "s=##part_chip_prod##=${part_chip_prod/\/dev\//\/dev\/block\/}=g"   $fstab_tmp_file
    debug_print "$(cat $fstab_tmp_file)"
}

tmp_img_dir=
temp_images_mnt_dir=
gen_fstab_file
if [[ $is_local_pc -eq 0 ]]; then
    while [[ $ret -ne 0 ]];
    do
        ping $target_pc_ip -w 1 -q
        ret=$?
        if [[ $ret -ne 0 ]]; then
            echo "board can not is reachable"
        fi
    done;

    echo "begin to install ..."
    tmp_img_dir=$(ssh root@$target_pc_ip "mktemp -d")
    check_temp_dir $tmp_img_dir
    scp ${image_path}/system.img root@$target_pc_ip:$tmp_img_dir
    scp ${image_path}/vendor.img root@$target_pc_ip:$tmp_img_dir
    scp ${image_path}/ramdisk.img root@$target_pc_ip:/boot/ramdisk.img
    scp ${image_path}/bzImage root@$target_pc_ip:/boot/
    scp ${image_path}/sys_prod.img root@$target_pc_ip:$tmp_img_dir
    scp ${image_path}/chip_prod.img root@$target_pc_ip:$tmp_img_dir
else
    tmp_img_dir=$image_path
fi

if [[ $is_local_pc -eq 0 ]]; then
    src=$(ssh root@$target_pc_ip "mktemp -d")
    dst=$(ssh root@$target_pc_ip "mktemp -d")
    check_temp_dir $src
    check_temp_dir $dst
    ssh root@$target_pc_ip "mount $tmp_img_dir/system.img $src && mount ${part_system} $dst && rm -rf $dst/* && cp -ardf $src/* $dst/ && umount $src && umount $dst"
    ssh root@$target_pc_ip "mount $tmp_img_dir/vendor.img $src && mount ${part_vendor} $dst && rm -rf $dst/* && cp -ardf $src/* $dst/ && umount $src && umount $dst"
    ssh root@$target_pc_ip "mount $tmp_img_dir/sys_prod.img $src && mount ${part_sys_prod} $dst && rm -rf $dst/* && cp -ardf $src/* $dst/ && umount $src && umount $dst"
    ssh root@$target_pc_ip "mount $tmp_img_dir/chip_prod.img $src && mount ${part_chip_prod} $dst && rm -rf $dst/* && cp -ardf $src/* $dst/ && umount $src && umount $dst"
    ssh root@$target_pc_ip "mount ${part_data} $dst && rm -rf $dst/* && umount $dst"

    ssh root@$target_pc_ip "mount ${part_vendor} $dst"
    scp $fstab_tmp_file root@$target_pc_ip:$dst/etc/fstab.jiuxiao
    ssh root@$target_pc_ip "umount $dst && rm -fr $src && rm -fr $dst && rm -fr $tmp_img_dir"
    if [[ $flag_reboot -eq 1 ]]; then
        echo "reboot $target_pc_ip"
        ssh root@$target_pc_ip "reboot"
    fi
else
    src=$(mktemp -d)
    dst=$(mktemp -d)
    echo "install ramdisk image..."
    cp -fra "${image_path}/ramdisk.img" /boot/ramdisk.img
    echo "install bzImage image..."
    cp -fra "${image_path}/bzImage" /boot/
    echo "install system image..."
    mount "$tmp_img_dir/system.img" $src && mount ${part_system} $dst && rm -rf $dst/* && cp -ardf $src/* $dst/ && umount $src && umount $dst
    echo "install vendor image..."
    mount "$tmp_img_dir/vendor.img" $src && mount ${part_vendor} $dst && rm -rf $dst/* && cp -ardf $src/* $dst/ && umount $src && umount $dst
    echo "install sys_prod image..."
    mount "$tmp_img_dir/sys_prod.img" $src && mount ${part_sys_prod} $dst && rm -rf $dst/* && cp -ardf $src/* $dst/ && umount $src && umount $dst
    echo "install chip_prod image..."
    mount "$tmp_img_dir/chip_prod.img" $src && mount ${part_chip_prod} $dst && rm -rf $dst/* && cp -ardf $src/* $dst/ && umount $src && umount $dst
    echo "reset data partion..."
    mount "${part_data}" $dst && rm -rf $dst/* && umount $dst

    echo "update fstab..."
    mount ${part_vendor} $dst
    cp -fra $fstab_tmp_file $dst/etc/fstab.jiuxiao
    umount $dst && rm -fr $src && rm -fr $dst
    if [[ $flag_reboot -eq 1 ]]; then
        echo "reboot $target_pc_ip"
        reboot
    fi
fi

if [[ $is_unc_path -eq 1 ]]; then
    umount $temp_image_path_for_unc_mount
    rm -fr $temp_image_path_for_unc_mount
fi

echo "=====x86 install  successful=====!"
